import os
import lxml.etree as ET
import cv2
import matplotlib.pyplot as plt
import albumentations as A

from params import *
from model import *

image = "Base_100001"

# === 读取 XML 标注 ===
xml_path = os.path.join(annotation_dir, f"{image}.xml")
ann_tree = ET.parse(xml_path)
root = ann_tree.getroot()

# === 加载图像 ===
img_path = os.path.join(image_dir, f"{image}.{image_type}")
img = cv2.imread(img_path)

# === 解析标注框 ===
bboxes = []
labels = []
for obj in root.findall("object"):
    bbox = obj.find("bndbox")
    xmin = int(bbox.find("xmin").text)
    ymin = int(bbox.find("ymin").text)
    xmax = int(bbox.find("xmax").text)
    ymax = int(bbox.find("ymax").text)
    label = obj.find("name").text
    bboxes.append([xmin, ymin, xmax, ymax])
    labels.append(label)

# === albumentations 增强管道 ===
transform = A.Compose(
    [
        A.OneOf(
            [
                A.ElasticTransform(alpha=120, sigma=120 * 0.05, alpha_affine=20, p=1.0),
                A.GridDistortion(num_steps=5, distort_limit=0.3, p=1.0),
                A.OpticalDistortion(distort_limit=0.5, shift_limit=0.05, p=1.0),
                A.Perspective(scale=(0.05, 0.1), p=1.0),
            ],
            p=1.0,
        )
    ],
    bbox_params=A.BboxParams(format="pascal_voc", label_fields=["category_ids"]),
)

# === 执行增强 ===
aug = transform(image=img, bboxes=bboxes, category_ids=labels)
aug_img = aug["image"]
aug_bboxes = aug["bboxes"]
aug_labels = aug["category_ids"]

# === 可视化 ===
for box in aug_bboxes:
    x1, y1, x2, y2 = map(int, box)
    _image = aug_img
    cv2.rectangle(_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
plt.imshow(cv2.cvtColor(aug_img, cv2.COLOR_BGR2RGB))
plt.title("Albumentations Augmented")
plt.axis("off")
plt.show()

# === 保存图像 ===
save_name = f"{image}_distorted"
cv2.imwrite(os.path.join(image_dir, f"{save_name}.png"), aug_img)

# === 保存新的 XML ===
new_root = ET.Element("annotation")
ET.SubElement(new_root, "filename").text = f"{save_name}.png"
size = ET.SubElement(new_root, "size")
ET.SubElement(size, "width").text = str(aug_img.shape[1])
ET.SubElement(size, "height").text = str(aug_img.shape[0])
ET.SubElement(size, "depth").text = str(aug_img.shape[2])

for label, box in zip(aug_labels, aug_bboxes):
    obj = ET.SubElement(new_root, "object")
    ET.SubElement(obj, "name").text = label
    bndbox = ET.SubElement(obj, "bndbox")
    x1, y1, x2, y2 = map(int, box)
    ET.SubElement(bndbox, "xmin").text = str(x1)
    ET.SubElement(bndbox, "ymin").text = str(y1)
    ET.SubElement(bndbox, "xmax").text = str(x2)
    ET.SubElement(bndbox, "ymax").text = str(y2)

new_tree = ET.ElementTree(new_root)
new_tree.write(os.path.join(annotation_dir, f"{save_name}.xml"), pretty_print=True)
